from AStock.ASQuery import (
    ASQuery_stock_info,
    ASQuery_valuation,
    ASQuery_stock_pool
)
import argparse
from datetime import datetime
from AStock.ASSetting import output_path
import os


def get_stock_codes_from_file(stock_codes_file):
    codes = []
    with open(stock_codes_file, 'rt', encoding='utf-8') as f:
        for row in f:
            code = row.split()[0]
            if code:
                codes.append(code)
    return codes


def _main_(args):
    stock_pool = args.stock_pool
    if args.file:
        output_file = args.file  # 导出文件名
    else:
        if stock_pool:
            output_file = os.path.join(
                output_path,
                'pool_valuation_{}.xlsx'.format(datetime.now().strftime('%Y%m%d'))
            )
        else:
            output_file = os.path.join(
                output_path,
                'valuation_{}.xlsx'.format(datetime.now().strftime('%Y%m%d'))
            )
    stock_codes = args.stock_codes  # list，需要估值的股票代码
    # str, 需要估值的股票代码文件，每行为 股票代码和股票名称，空格分隔
    stock_codes_file = args.stock_codes_file

    if stock_codes_file:
        stock_codes = get_stock_codes_from_file(stock_codes_file)

    if stock_pool:
        df_pool = ASQuery_stock_pool()
        stock_codes = df_pool['code'].tolist()
        df_pool = df_pool.drop('name', axis=1)

    # if stock_codes:
    #     print(stock_codes)

    # 获取股票信息
    df_stock_info = ASQuery_stock_info(code=stock_codes)
    # print(df_stock_info.tail(20))
    # 获取估值
    df_pe_percentage = ASQuery_valuation(code=stock_codes)
    # print(df_pe_percentage.tail(20))

    # join
    df = df_stock_info.join(df_pe_percentage.set_index('code'), on='code', how='left')
    if stock_pool:
        df = df.join(df_pool.set_index('code'), on='code', how='left')

    # 计算PE
    df = df.assign(
        thisYearPrice30=df.predictThisYearEps * df.pe30,
        thisYearPrice50=df.predictThisYearEps * df.pe50,
        thisYearPrice70=df.predictThisYearEps * df.pe70,
        nextYearPrice30=df.predictNextYearEps * df.pe30,
        nextYearPrice50=df.predictNextYearEps * df.pe50,
        nextYearPrice70=df.predictNextYearEps * df.pe70
    )

    df = df.assign(
        thisYearPrice30D=(df.close - df.thisYearPrice30) / df.thisYearPrice30,
        thisYearPrice50D=(df.close - df.thisYearPrice50) / df.thisYearPrice50,
        thisYearPrice70D=(df.close - df.thisYearPrice70) / df.thisYearPrice70,
        nextYearPrice30D=(df.close - df.nextYearPrice30) / df.nextYearPrice30,
        nextYearPrice50D=(df.close - df.nextYearPrice50) / df.nextYearPrice50,
        nextYearPrice70D=(df.close - df.nextYearPrice70) / df.nextYearPrice70
    )
    # 计算PEG


    # print(df.columns)
    columns = [
        'code',
        'name',
        'close',
        'close_date',
        #'predictThisYearEps',
        #'predictNextYearEps',
        'predictThisYearPe',
        'predictNextYearPe',
        'predictThisYearPeg',
        'predictNextYearPeg',
        'pe30',
        'pe50',
        'pe70',
        #'thisYearPrice30',
        #'thisYearPrice50',
        #'thisYearPrice70',
        #'nextYearPrice30',
        #'nextYearPrice50',
        #'nextYearPrice70',
        'thisYearPrice30D',
        'thisYearPrice50D',
        'thisYearPrice70D',
        'nextYearPrice30D',
        'nextYearPrice50D',
        'nextYearPrice70D',
        'recentMonthReportCount',
        'industry',
        'area'
    ]
    if stock_pool:
        columns.append('groups')

    df = df.round(2)
    if output_file.endswith('.csv'):
        df.to_csv(
            output_file,
            index=False,
            line_terminator='\n',
            encoding='utf-8',
            columns=columns
        )
    elif output_file.endswith('.xlsx') or output_file.endswith('.xls'):
        df.to_excel(
            output_file,
            index=False,
            encoding='utf-8',
            columns=columns
        )
    print('saved to {}'.format(output_file))


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='export valuation')
    argparser.add_argument('-s', '--stock-codes', nargs='+',
                           help='stock codes for valuation')
    argparser.add_argument('-c', '--stock-codes-file', required=False,
                           help='stock codes file for valuation')
    argparser.add_argument('-f', '--file', required=False,
                           help='full path file name to export to')
    argparser.add_argument('--stock-pool', required=False, action='store_true',
                           help='export valuation of stocks from stock pool')
    args = argparser.parse_args()
    _main_(args)

